from time import perf_counter

from sahi.model import Yolov5DetectionModel
from sahi.utils.cv import read_image
from sahi.predict import get_prediction, get_sliced_prediction
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

yolov5_model_path = '/home/hao/Downloads/exp6/weights/best.pt'

detection_model = Yolov5DetectionModel(
    model_path=yolov5_model_path,
    confidence_threshold=0.3,
    device="cpu",  # or 'cuda:0'
)

img = "/home/hao/Code/python/A10/datasets/board96/images_augmented/train/resized59.jpg"

start_time = perf_counter()
result = get_prediction(read_image(img), detection_model)
print(perf_counter() - start_time)

# result.export_visuals(export_dir="yolo_out/")

start_time = perf_counter()
result = get_sliced_prediction(
    img,
    detection_model,
    slice_height=256,
    slice_width=256,
    overlap_height_ratio=0.2,
    overlap_width_ratio=0.2
)
print(perf_counter() - start_time)

# result.export_visuals(export_dir="sahi_out/")

object_prediction_list = result.object_prediction_list

object_prediction_list[0]

result.to_coco_annotations()[:3]

result.to_coco_predictions(image_id=1)[:3]
